import time
import threading


class SnowflakeSingleton:
    """
    单例模式的雪花算法生成唯一ID工具类
    """
    _instance = None
    _lock = threading.Lock()  # 单例模式的线程锁

    # 起始时间戳 (2020-01-01 00:00:00)
    START_TIMESTAMP = 1577808000000

    # 各部分位数
    SEQUENCE_BITS = 12  # 序列号位数
    WORKER_ID_BITS = 10  # 工作节点ID位数

    # 最大取值计算
    MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)  # 1023
    MAX_SEQUENCE = -1 ^ (-1 << SEQUENCE_BITS)  # 4095

    # 移位偏移量计算
    WORKER_ID_SHIFT = SEQUENCE_BITS  # 12
    TIMESTAMP_LEFT_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS  # 22

    def __new__(cls, worker_id=None):
        """单例模式实现，确保只创建一个实例"""
        with cls._lock:
            # 如果实例不存在，则创建
            if not cls._instance:
                # 首次创建时必须提供worker_id
                if worker_id is None:
                    raise ValueError("首次初始化必须提供worker_id参数")
                cls._instance = super().__new__(cls)
                # 初始化实例变量
                cls._instance.worker_id = worker_id
                cls._instance.last_timestamp = -1  # 上次生成ID的时间戳
                cls._instance.sequence = 0  # 序列号
                cls._instance.id_lock = threading.Lock()  # 生成ID的线程锁
            # 如果已经存在实例，再次传入的worker_id必须与初始一致
            elif worker_id is not None and worker_id != cls._instance.worker_id:
                raise ValueError(f"雪花算法实例已初始化，worker_id必须为{cls._instance.worker_id}")
            return cls._instance

    @classmethod
    def get_instance(cls, worker_id=None):
        """获取单例实例的便捷方法"""
        return cls(worker_id)

    def _gen_timestamp(self):
        """生成当前时间戳（毫秒）"""
        return int(time.time() * 1000)

    def get_id(self):
        """生成下一个唯一ID"""
        with self.id_lock:  # 确保线程安全
            timestamp = self._gen_timestamp()

            # 处理系统时钟回退
            if timestamp < self.last_timestamp:
                raise RuntimeError(
                    f"系统时钟回退，拒绝生成ID。时间差: {self.last_timestamp - timestamp}毫秒"
                )

            # 处理同一毫秒内的序列号
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                # 序列号达到最大值，等待下一毫秒
                if self.sequence == 0:
                    timestamp = self._til_next_millis(self.last_timestamp)
            else:
                # 不同时间戳，序列号重置为0
                self.sequence = 0

            # 更新上次生成ID的时间戳
            self.last_timestamp = timestamp

            # 组合生成ID
            snowflake_id = (
                    ((timestamp - self.START_TIMESTAMP) << self.TIMESTAMP_LEFT_SHIFT) |
                    (self.worker_id << self.WORKER_ID_SHIFT) |
                    self.sequence
            )

            return snowflake_id

    def _til_next_millis(self, last_timestamp):
        """阻塞到下一个毫秒，直到获得新的时间戳"""
        timestamp = self._gen_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._gen_timestamp()
        return timestamp


# 调用方法示例
if __name__ == "__main__":
    # 初始化单例（首次调用必须提供worker_id）
    snowflake = SnowflakeSingleton.get_instance(worker_id=1)

    # 生成10个ID并打印
    for _ in range(10):
        print(f"生成的ID: {snowflake.get_id()}")

    # 在其他地方获取实例（无需再次提供worker_id）
    another_instance = SnowflakeSingleton.get_instance()
    print(f"\n验证单例: {snowflake is another_instance}")  # 应输出True

    # 多线程测试
    def generate_ids(count):
        ids = []
        for _ in range(count):
            ids.append(SnowflakeSingleton.get_instance().get_id())
        return ids

    # 创建多个线程同时生成ID
    threads = []
    for i in range(5):
        t = threading.Thread(target=generate_ids, args=(1000,))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()

    print("\n多线程ID生成测试完成")